"""
核反应堆子通道分析程序 - 基于两流体模型

本程序实现了一个简化的核反应堆子通道分析工具，用于模拟和分析燃料棒束中冷却剂的流动特性。
采用两流体模型处理气液两相流动，考虑了轴向和横向的质量、动量传递。

主要功能：
1. 模拟子通道内的两相流动
2. 计算空泡率、压力分布
3. 分析液相和气相的速度场
4. 考虑子通道间的横向混合

物理模型：
- 两流体模型：分别处理气相和液相
- 考虑相间作用力
- 包含重力和压力梯度效应
- 简化的横向混合模型

数值方法：
- 采用显式时间推进格式
- 交错网格离散
- 迭代求解直至收敛

作者：[作者名称]
日期：[日期]
"""

import numpy as np
from CoolProp.CoolProp import PropsSI
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import os
import csv
from matplotlib.animation import FuncAnimation
import multiprocessing as mp
from functools import partial
from scipy.sparse import csr_matrix, lil_matrix
from scipy.sparse.linalg import gmres
import numpy.linalg as la
from numba import jit, njit, prange
import numba as nb
import json
import traceback

plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文显示
plt.rcParams['axes.unicode_minus'] = False  # 设置负号显示


class SubchannelAnalysis:
    """
    子通道分析的主类，实了两流体模型的求解
    
    主要属性：
    - n_axial: 轴向节点数
    - n_channels: 子通道数量
    - channel_connections: 子通道连接关系
    - connection_coefficients: 子通道间的连接系数
    
    主要方法：
    - solve_mass_momentum: 求解质量和动量守恒方程
    - calculate_cross_flow: 计算横向流动
    - get_fluid_properties: 获取流体物性
    """
    
    def __init__(self, n_axial=1500):
        """
        初始化子通道分析器
        
        参数：
        n_axial: int, 轴向节点数
        """
        # ====== 网格参数 ======
        self.n_axial = n_axial                    # 轴向节点数
        self.channel_length = 4.0                  # 通道总长度 (m)
        self.dz = self.channel_length / n_axial    # 轴向网格间距 (m)
        
        # ====== 几何参数 ======
        self.channel_diameter = 0.015              # 子通道等效直径 (m)
        self.channel_area = np.pi * (self.channel_diameter/2)**2  # 子通道横截面积 (m²)
        
        # ====== 子通道拓扑结构 ======
        # 定义子通道连接关系（根据图2-2）
        self.channel_connections = {
            1: [2],           # 1号子通道与2号相邻
            2: [1, 3, 4],     # 2号子通道与1、3、4号相邻
            3: [2, 4],        # 3号子通道与2、4号相邻
            4: [2, 3, 5, 6],  # 4号子通道与2、3、5、6号相邻
            5: [4, 6],        # 5号子通道与4、6号相邻
            6: [4, 5]         # 6号子通道与4、5号相邻
        }
        
        # 定义子通道类型
        self.channel_types = {
            1: 'corner',      # 角落子通道
            2: 'side',        # 侧边子通道
            3: 'side',        # 侧边子通道
            4: 'center',      # 中心子通道
            5: 'side',        # 侧边子通道
            6: 'corner'       # 角落子通道
        }
        
        # 定义加热边界（燃料棒位置）
        self.heated_boundaries = {
            1: [(2, 0.25)],           # 1号子通道被1/4根燃料棒加热
            2: [(2, 0.5), (1, 0.5)],  # 2号子通道被两根燃料棒各加热1/2
            3: [(2, 0.5)],            # 3号子通道被1/2根燃料棒加热
            4: [(4, 1.0)],            # 4号子通道被1根完整燃料棒加热
            5: [(2, 0.5)],            # 5号子通道被1/2根燃料棒加热
            6: [(1, 0.25)]            # 6号子通道被1/4根燃料棒加热
        }
        
        # 定义子通道几何参数
        self.channel_geometry = {
            'corner': {
                'flow_area': 0.25 * np.pi * self.channel_diameter**2,
                'wetted_perimeter': 0.5 * np.pi * self.channel_diameter,
                'heated_perimeter': 0.25 * np.pi * self.channel_diameter
            },
            'side': {
                'flow_area': 0.5 * np.pi * self.channel_diameter**2,
                'wetted_perimeter': np.pi * self.channel_diameter,
                'heated_perimeter': 0.5 * np.pi * self.channel_diameter
            },
            'center': {
                'flow_area': np.pi * self.channel_diameter**2,
                'wetted_perimeter': 2 * np.pi * self.channel_diameter,
                'heated_perimeter': np.pi * self.channel_diameter
            }
        }
        
        # ====== 入口工况参数 ======
        self.inlet_pressure = 15.5e6  # 入口压力 (Pa)
        self.inlet_temp = 300          # 入口温度 (K)
        self.mass_flux = 1500          # 量流量 (kg/m²s)
        
        # ====== 场变量初始化 ======
        self.n_channels = 6           # 子通道总数
        # 创建各个物理量的二维数组 (轴向节点 × 子通道数)
        self.void_fraction = np.zeros((n_axial, self.n_channels))    # 空泡率 (-)

        # 设置入口处的空泡率
        self.void_fraction[0, :] = 0.1  # 假设入口处有10%的空泡率

        # 初始化液相和气相速度
        self.liquid_velocity = np.zeros((n_axial, self.n_channels))  # 液相速度 (m/s)
        self.vapor_velocity = np.zeros((n_axial, self.n_channels))   # 气相速度 (m/s)

        # 根据质量流量和液相密度计算入口液相速度
        rho_l = self.get_fluid_properties(self.inlet_pressure, self.inlet_temp, 'liquid')['density']
        self.liquid_velocity[0, :] = self.mass_flux / rho_l * (1 - self.void_fraction[0, :])  # 入口液相速度

        # 设置入口气相速度，假设气相速度为液相速度的10%
        self.vapor_velocity[0, :] = self.liquid_velocity[0, :] * 0.1  # 入口气相速度

        # 压力和温度场初始化为入口条件
        self.pressure = np.full((n_axial, self.n_channels), self.inlet_pressure)  # 压力场 (Pa)
        self.temperature = np.full((n_axial, self.n_channels), self.inlet_temp)   # 温度场 (K)
        self.liquid_enthalpy = np.zeros((n_axial, self.n_channels))  # 液相焓 (J/kg)
        self.vapor_enthalpy = np.zeros((n_axial, self.n_channels))   # 气相焓 (J/kg)
        
        # 添加两流体模型所需的物理量
        self.alpha_l = np.zeros((n_axial, self.n_channels))  # 液相分数
        self.alpha_v = np.zeros((n_axial, self.n_channels))  # 气相分数
        
        # 速度场
        self.u_l = np.zeros((n_axial, self.n_channels))  # 液相轴向速度
        self.u_v = np.zeros((n_axial, self.n_channels))  # 气相轴向速度
        self.v_l = np.zeros((n_axial, self.n_channels))  # 液相横向速度
        self.v_v = np.zeros((n_axial, self.n_channels))  # 气相横向速度
        self.w_l = np.zeros((n_axial, self.n_channels))  # 液相z方向速度
        self.w_v = np.zeros((n_axial, self.n_channels))  # 气相z方向速度
        
        # 焓和温度场
        self.h_l = np.zeros((n_axial, self.n_channels))  # 液相焓
        self.h_v = np.zeros((n_axial, self.n_channels))  # 气相焓
        self.T_l = np.zeros((n_axial, self.n_channels))  # 液相温度
        self.T_v = np.zeros((n_axial, self.n_channels))  # 气相温度
        self.T_c = np.zeros((n_axial, self.n_channels))  # 包壳温度
        
        # 相变和传热相关参数
        self.gamma = np.zeros((n_axial, self.n_channels))  # 相变质量流量
        self.q_wl = np.zeros((n_axial, self.n_channels))  # 液相���面热流密度
        self.q_wv = np.zeros((n_axial, self.n_channels))  # 气相壁面热流密度
        
        # 添加瞬态计算相关的参数
        self.total_time = 5       # 总计算时间 (s)
        self.time_step = 0.001        # 时间步长 (s)
        self.current_time = 0.0       # 当前时间 (s)
        self.max_iterations = 1000    # 每个时间步内的最大迭代次数
        self.convergence_criteria = 1e-6  # 收敛判据
        
        # 保存上一时间步的场值，用于时间导数计算
        self.prev_void_fraction = np.zeros((n_axial, self.n_channels))
        self.prev_liquid_velocity = np.zeros((n_axial, self.n_channels))
        self.prev_vapor_velocity = np.zeros((n_axial, self.n_channels))
        self.prev_pressure = np.zeros((n_axial, self.n_channels))
        self.prev_temperature = np.zeros((n_axial, self.n_channels))
        
        # 初始化物理参数
        self._initialize_physical_parameters()
        
        # 添加结果保存相关的参数
        self.save_interval = 100  # 每隔多少时间步保存一次结果
        self.results_folder = 'transient_results'  # 果保存文件夹名称
        
        # 初始化所有场变量
        self.void_fraction = np.zeros((n_axial, self.n_channels))
        self.pressure = np.full((n_axial, self.n_channels), self.inlet_pressure)
        self.temperature = np.full((n_axial, self.n_channels), self.inlet_temp)
        self.liquid_velocity = np.zeros((n_axial, self.n_channels))
        self.vapor_velocity = np.zeros((n_axial, self.n_channels))
        self.T_l = np.full((n_axial, self.n_channels), self.inlet_temp)
        self.T_v = np.full((n_axial, self.n_channels), self.inlet_temp)
        self.T_c = np.full((n_axial, self.n_channels), self.inlet_temp)
        self.q_wl = np.zeros((n_axial, self.n_channels))
        self.q_wv = np.zeros((n_axial, self.n_channels))
        self.gamma = np.zeros((n_axial, self.n_channels))
        
        # 初始化历史数据字典
        self.history = {
            # 时间相关
            'time': [],
            'time_steps': [],
            
            # 场变量历史
            'void_fraction_history': [],
            'pressure_history': [],
            'temperature_history': [],
            'liquid_velocity_history': [],
            'vapor_velocity_history': [],
            'liquid_temp_history': [],
            'vapor_temp_history': [],
            'wall_temp_history': [],
            'heat_flux_liquid_history': [],
            'heat_flux_vapor_history': [],
            'phase_change_rate_history': [],
            
            # 统计量历史
            'max_void_fraction': [],
            'avg_pressure': [],
            'max_temperature': [],
            'max_liquid_velocity': [],
            'max_vapor_velocity': [],
            'residuals': [],
            'iterations': []
        }
        
        # 创建结果保存文件夹
        self.results_folder = 'transient_results'
        if not os.path.exists(self.results_folder):
            os.makedirs(self.results_folder)
            
    def _initialize_physical_parameters(self):
        """初始化物理参数"""
        # 设置入口条件
        self.alpha_l[0,:] = 1.0  # 入口处全为液相
        self.alpha_v[0,:] = 0.0
        
        # 根据质量流量设置入口速度
        rho_l = self.get_fluid_properties(self.inlet_pressure, self.inlet_temp, 'liquid')['density']
        self.u_l[0,:] = self.mass_flux / rho_l
        self.u_v[0,:] = 0.0
        
        # 设置入口温度和焓
        self.T_l[0,:] = self.inlet_temp
        self.T_v[0,:] = self.inlet_temp
        self.h_l[0,:] = self.get_fluid_properties(self.inlet_pressure, self.inlet_temp, 'liquid')['enthalpy']
        self.h_v[0,:] = self.get_fluid_properties(self.inlet_pressure, self.inlet_temp, 'vapor')['enthalpy']
    
    def get_fluid_properties(self, p, t, phase):
        """
        获取指定状态点的流体物性
        
        参数：
        p: float, 压力 (Pa)
        t: float, 温度 (K)
        phase: str, 相态 ('liquid' 或 'vapor')
        
        返回：
        dict, 包含密度、粘度和焓值的字典
        """
        try:
            # 确保参数在CoolProp的有效范围内
            p = max(611.655, min(p, 2.2064e7))  # 水的三相点到临界点压力范围
            t = max(273.16, min(t, 647.096))    # 水的三相点到临界点温度范围
            
            if phase == 'liquid':
                # 计算液相物性
                return {
                    'density': PropsSI('D', 'P', p, 'T', t, 'Water'),    # 密度 (kg/m³)
                    'viscosity': PropsSI('V', 'P', p, 'T', t, 'Water'),  # 动力粘度 (Pa·s)
                    'enthalpy': PropsSI('H', 'P', p, 'T', t, 'Water')    # 比焓 (J/kg)
                }
            else:  # vapor
                # 计算气相物性
                return {
                    'density': PropsSI('D', 'P', p, 'Q', 1, 'Water'),    # 密度
                    'viscosity': PropsSI('V', 'P', p, 'Q', 1, 'Water'),  # 动力粘度
                    'enthalpy': PropsSI('H', 'P', p, 'Q', 1, 'Water')    # 比焓
                }
        except Exception as e:
            print(f"警：物性计算失败 (P={p}, T={t})")
            print(f"错误信息: {str(e)}")
            return None
            
    def calculate_cross_flow(self, i, j, k):
        """
        计算子通道间的横向流动
        
        参数：
        i: int, 轴向位置索引
        j: int, 当前子通道编号
        k: int, 相邻子通道编号
        
        返回：
        float, 横向质量流量 (kg/s)
        """
        # 检查是否存在连接
        if (j, k) not in self.connection_coefficients:
            return 0.0
        
        # 获取连接系数    
        s_jk = self.connection_coefficients[(j, k)]
        
        # 计算压力差驱动的横向流动
        dp = self.pressure[i, j-1] - self.pressure[i, k-1]  # 压力差
        cross_flow = s_jk * dp * 1e-3  # 横向流与压力差成正比
        
        return cross_flow
        
    def solve_mass_momentum(self):
        """
        求解质量和动量守恒方程
        
        采用显式时间推进格式解：
        1. 质量守恒方程用于更新空泡率
        2. 动量方程用于更新速度场
        3. 考虑横向流动的影响
        """
        # 数值求解参数
        dt = 1e-4          # 时间步长 (s)
        max_iterations = 100000  # 最大迭代次数
        tolerance = 1e-6    # 收敛判据
        
        # 迭代求解
        for n in range(max_iterations):
            max_error = 0.0
            
            # 保存上一时间步的值用于计算误差
            old_void = np.copy(self.void_fraction)
            old_vl = np.copy(self.liquid_velocity)
            old_vv = np.copy(self.vapor_velocity)
            
            # 空间推进
            for i in range(1, self.n_axial):
                for j in range(self.n_channels):
                    # 获取当前位置的流体物性
                    props_l = self.get_fluid_properties(
                        float(self.pressure[i,j]), 
                        float(self.temperature[i,j]), 
                        'liquid'
                    )
                    props_v = self.get_fluid_properties(
                        float(self.pressure[i,j]), 
                        float(self.temperature[i,j]), 
                        'vapor'
                    )
                    
                    if props_l is None or props_v is None:
                        continue
                    
                    # 提取物性值    
                    rho_l = props_l['density']    # 液相密度
                    rho_v = props_v['density']    # 气相密度
                    
                    # 计算横向流动贡献
                    cross_flows = []
                    for k in self.channel_connections[j+1]:
                        w_jk = self.calculate_cross_flow(i, j+1, k)
                        cross_flows.append(w_jk)
                    
                    # 提取当前节点的场变量值
                    alpha = self.void_fraction[i-1,j]     # 空泡率
                    vl = self.liquid_velocity[i-1,j]      # 液相速度
                    vv = self.vapor_velocity[i-1,j]       # 气相速度
                    
                    # 求解空泡率输运方程（显式格式）
                    new_alpha = alpha + dt * (
                        -vv * (alpha - self.void_fraction[i-1,j]) / self.dz  # 对流项
                        + sum(cross_flows) / (rho_v * self.channel_area)      # 横向流动项
                    )
                    new_alpha = np.clip(new_alpha, 0, 0.9999)  # 限制在物理范围内
                    
                    # 求解液相动量方程
                    new_vl = vl + dt * (
                        -vl * (vl - self.liquid_velocity[i-1,j]) / self.dz    # 对流项
                        - (1/rho_l) * (self.pressure[i,j] - self.pressure[i-1,j]) / self.dz  # 压力梯度
                        - 9.81  # 重力项
                    )
                    
                    # 求解气相动量方程
                    new_vv = vv + dt * (
                        -vv * (vv - self.vapor_velocity[i-1,j]) / self.dz     # 对流项
                        - (1/rho_v) * (self.pressure[i,j] - self.pressure[i-1,j]) / self.dz  # 压力梯度
                        - 9.81  # 重力项
                    )
                    
                    # 更新场变量
                    self.void_fraction[i,j] = new_alpha
                    self.liquid_velocity[i,j] = new_vl
                    self.vapor_velocity[i,j] = new_vv
                    
                    # 计算最大误差
                    error = max(
                        abs(new_alpha - old_void[i,j]),
                        abs(new_vl - old_vl[i,j]),
                        abs(new_vv - old_vv[i,j])
                    )
                    max_error = max(max_error, error)
            
            # 检查收敛性
            if max_error < tolerance:
                print(f"计在第{n}次迭代后收敛")
                break
                
            if n == max_iterations - 1:
                print("警告：计算未收敛")
    
    def run_simulation(self):
        """运行完整的模拟计算"""
        print("开始数值计算...")
        
        time = 0
        step = 0
        
        while time < self.total_time:
            print(f"当前时间：{time:.3f} s")
            
            # 使用自适应时间步长求解
            success = self.solve_two_fluid_equations()
            
            if not success:
                print(f"警告：在时间 {time:.3f} s 处计算不稳定")
                self.time_step *= 0.1
                if self.time_step < 1e-8:
                    print("错误：时间步长过小，计算终止")
                    break
                continue
            
            # 更新时间
            time += self.time_step
            step += 1
            
            # 保存结果
            if step % self.save_interval == 0:
                self.save_timestep_results(step)
                
        print("计算完成")
    
    def plot_results(self):
        """绘制计算结果的可视化图，增加错误检查和空值处理"""
        try:
            # 创建网格点
            x = np.arange(self.n_axial)
            y = np.arange(self.n_channels)
            X, Y = np.meshgrid(x, y)
            
            # 检查数据是否为空
            if np.all(self.void_fraction == 0) or np.all(self.pressure == 0):
                print("警告：检测到空数据，请确保计算已完成")
                return
                
            # 创建两个图形窗口
            fig1 = plt.figure(figsize=(20, 15))
            fig2 = plt.figure(figsize=(20, 15))
            
            # 第一个图形窗口：液相参数
            axes1 = []
            for i in range(6):
                ax = fig1.add_subplot(2, 3, i+1, projection='3d')
                axes1.append(ax)
                
            # 第二个图形窗口：气相参数
            axes2 = []
            for i in range(6):
                ax = fig2.add_subplot(2, 3, i+1, projection='3d')
                axes2.append(ax)
                
            # 绘制液相参数
            plots = [
                (self.alpha_l, '液相分数分布'),
                (self.u_l, '液相轴向速度分布 (m/s)'),
                (self.v_l, '液相横向速度分布 (m/s)'),
                (self.T_l, '液相温度分布 (K)'),
                (self.pressure, '压力分布 (Pa)'),
                (self.q_wl if hasattr(self, 'q_wl') else np.zeros_like(self.pressure), 
                 '液相壁面热流密度 (W/m²)')
            ]
            
            for ax, (data, title) in zip(axes1, plots):
                surf = ax.plot_surface(X, Y, data.T, cmap=cm.viridis)
                ax.set_title(title)
                ax.set_xlabel('轴向位置')
                ax.set_ylabel('子通道编号')
                fig1.colorbar(surf, ax=ax, shrink=0.5, aspect=5)
                
            # 绘制气相参数
            plots = [
                (self.alpha_v, '气相分数分布'),
                (self.u_v, '气相轴向速度分布 (m/s)'),
                (self.v_v, '气相横向速度分布 (m/s)'),
                (self.T_v, '气相温度分布 (K)'),
                (self.gamma if hasattr(self, 'gamma') else np.zeros_like(self.pressure), 
                 '相变率 (kg/m³·s)'),
                (self.q_wv if hasattr(self, 'q_wv') else np.zeros_like(self.pressure), 
                 '气相壁面热流密度 (W/m²)')
            ]
            
            for ax, (data, title) in zip(axes2, plots):
                surf = ax.plot_surface(X, Y, data.T, cmap=cm.viridis)
                ax.set_title(title)
                ax.set_xlabel('轴向位置')
                ax.set_ylabel('子通道编号')
                fig2.colorbar(surf, ax=ax, shrink=0.5, aspect=5)
                
            # 设置视角和网格
            for ax in axes1 + axes2:
                ax.view_init(elev=20, azim=45)
                ax.grid(True)
                
            plt.tight_layout()
            plt.show()
            
        except Exception as e:
            print(f"绘图时出错: {str(e)}")

    def solve_two_fluid_equations(self):
        """求解双流体模型方程组 - 添加自适应时间步长控制"""
        # 准备输入数据
        dx = dy = self.dz  # 假设网格等间距
        max_attempts = 5  # 最大重试次数
        
        for attempt in range(max_attempts):
            # 调用Numba优化的求解器
            results = solve_two_fluid_equations_numba(
                self.n_axial, self.n_channels,
                self.alpha_l, self.alpha_v,
                self.rho_l, self.rho_v,
                self.u_l, self.u_v,
                self.v_l, self.v_v,
                self.w_l, self.w_v,
                self.pressure, self.T_l, self.T_v, self.T_c,
                self.h_l, self.h_v,
                dx, dy, self.dz, self.time_step
            )
            
            # 解包结果
            if len(results) == 10:  # 正常返回
                (new_alpha_l, new_alpha_v, new_u_l, new_u_v,
                 new_v_l, new_v_v, new_h_l, new_h_v,
                 new_dt, needs_restart) = results
                
                if needs_restart:
                    print(f"时间步长调整为 {new_dt:.2e}，重新计算...")
                    self.time_step = new_dt
                    continue
                    
                # 更新时间步长
                self.time_step = new_dt
                
                # 更新场值
                self.alpha_l = new_alpha_l
                self.alpha_v = new_alpha_v
                self.u_l = new_u_l
                self.u_v = new_u_v
                self.v_l = new_v_l
                self.v_v = new_v_v
                self.h_l = new_h_l
                self.h_v = new_h_v
                
                print(f"计算成功，当前时间步长: {self.time_step:.2e}")
                return True
                
        print("警告：达到最大重试次数，计算可能不稳定")
        return False

    def calculate_phase_change(self, i, j):
        """计算相变质量流量"""
        T_c = self.T_c[i,j]  # 包壳温度
        T_sat = self.get_saturation_temperature(self.pressure[i,j])
        T_l = self.T_l[i,j]  # 液相温度
        
        # 过冷沸腾
        if T_l < T_sat and T_c > T_sat:
            h_l = self.calculate_subcooled_boiling_htc(T_c, T_sat, T_l, self.pressure[i,j])
            Q_l = np.pi * self.channel_diameter * h_l * (T_c - T_l)
            gamma_l_sc = Q_l / self.get_latent_heat(self.pressure[i,j])
            return -gamma_l_sc
            
        # 饱和核态沸腾
        elif T_l >= T_sat and T_c > T_sat:
            h_l = self.calculate_nucleate_boiling_htc(T_c, T_sat, self.pressure[i,j])
            Q_l = np.pi * self.channel_diameter * h_l * (T_c - T_sat)
            gamma_l_sh = Q_l / self.get_latent_heat(self.pressure[i,j])
            return gamma_l_sh
            
        return 0.0

    def calculate_turbulent_mixing(self, i, j, phase):
        """计算湍流混合和空泡漂移质量流量"""
        C_e = 0.01  # 涡流扩散率
        L_tur = self.channel_diameter  # 湍流混合长度
        
        mixing_mass = 0.0
        for k in self.channel_connections[j+1]:
            L_gap = self.get_gap_width(j+1, k)  # 获取间隙宽度
            
            if phase == 'liquid':
                rho_diff = self.get_fluid_properties(self.pressure[i,j], 
                    self.temperature[i,j], 'liquid')['density']
                alpha_diff = self.alpha_l[i,j] - self.alpha_l[i,k-1]
                alpha_eq_diff = self.calculate_equilibrium_void_fraction(i,j) - \
                               self.calculate_equilibrium_void_fraction(i,k-1)
            else:
                rho_diff = self.get_fluid_properties(self.pressure[i,j],
                    self.temperature[i,j], 'vapor')['density']
                alpha_diff = self.alpha_v[i,j] - self.alpha_v[i,k-1]
                alpha_eq_diff = (1-self.calculate_equilibrium_void_fraction(i,j)) - \
                               (1-self.calculate_equilibrium_void_fraction(i,k-1))
                
            mixing_mass += C_e * L_gap * L_tur * rho_diff * (alpha_diff - alpha_eq_diff)
            
        return mixing_mass

    def calculate_wall_friction(self, i, j, phase, direction='axial'):
        """计算壁面摩擦力
        
        参数:
        i, j: 网格位置
        phase: 'liquid'或'vapor'
        direction: 'axial'或'lateral'
        """
        if phase == 'liquid':
            rho = self.get_fluid_properties(self.pressure[i,j], self.T_l[i,j], 'liquid')['density']
            mu = self.get_fluid_properties(self.pressure[i,j], self.T_l[i,j], 'liquid')['viscosity']
            if direction == 'axial':
                v = self.u_l[i,j]
            else:
                v = self.v_l[i,j]
        else:
            rho = self.get_fluid_properties(self.pressure[i,j], self.T_v[i,j], 'vapor')['density']
            mu = self.get_fluid_properties(self.pressure[i,j], self.T_v[i,j], 'vapor')['viscosity']
            if direction == 'axial':
                v = self.u_v[i,j]
            else:
                v = self.v_v[i,j]
            
        # 计算Reynolds数
        Re = rho * abs(v) * self.channel_diameter / mu
        
        # 计算摩擦系数
        if Re < 2300:  # 层流
            f = 64 / Re
        else:  # 湍流
            f = 0.316 * Re**(-0.25)
            
        # 计算壁面摩擦力
        tau_w = 0.5 * f * rho * v * abs(v)
        
        # 考虑两相流修正系数
        if phase == 'liquid':
            phi_l = (1 - self.alpha_l[i,j])**(-1.75)  # 两相流修正系数
            tau_w *= phi_l
            
        return tau_w * np.pi * self.channel_diameter / self.channel_area

    def calculate_interfacial_force(self, i, j, direction='axial'):
        """计算界面剪切力 (eq.19)"""
        rho_l = self.get_fluid_properties(self.pressure[i,j], self.T_l[i,j], 'liquid')['density']
        
        if direction == 'axial':
            v_rel = self.u_v[i,j] - self.u_l[i,j]  # 相对速度
        else:
            v_rel = self.v_v[i,j] - self.v_l[i,j]
            
        # 计算界面面积密度
        a_i = 6 * self.alpha_v[i,j] / self.channel_diameter  # 假设球形气泡
        
        # 计算界面阻力系数
        C_i = 0.44  # 球形气泡的阻力系数
        
        # 计算界面剪切力
        tau_i = 0.5 * C_i * rho_l * v_rel * abs(v_rel) * a_i
        
        return tau_i

    def calculate_wall_heat_flux(self, i, j, phase):
        """计算壁面热流密度 (eq.6)"""
        if phase == 'liquid':
            # 根据沸腾状态确定换热系数
            T_sat = self.get_saturation_temperature(self.pressure[i,j])
            
            if self.T_l[i,j] < T_sat and self.T_c[i,j] > T_sat:
                # 过冷核态沸腾
                h = self.calculate_subcooled_boiling_htc(
                    self.T_c[i,j], T_sat, self.T_l[i,j], self.pressure[i,j])
                
            elif self.T_l[i,j] >= T_sat and self.T_c[i,j] > T_sat:
                # 饱和核态沸腾
                h = self.calculate_nucleate_boiling_htc(
                    self.T_c[i,j], T_sat, self.pressure[i,j])
                
            else:
                # 单相对流换热
                h = self.calculate_single_phase_htc(i, j, 'liquid')
                
            q = h * (self.T_c[i,j] - self.T_l[i,j])
        else:
            # 气相单相对流换热
            h = self.calculate_single_phase_htc(i, j, 'vapor')
            q = h * (self.T_c[i,j] - self.T_v[i,j])
            
        return q * np.pi * self.channel_diameter

    def calculate_subcooled_boiling_htc(self, T_w, T_sat, T_l, P):
        """计算过冷核态沸腾换热系数 (eq.8)"""
        F_gam = 0.5  # 气泡冷凝系数
        h = (1 - F_gam) * np.exp(2*P/1260 * (T_w-T_sat)/(2722*(T_w-T_l)))
        return h

    def calculate_nucleate_boiling_htc(self, T_w, T_sat, P):
        """计算饱和核态沸腾换热系数 (eq.9)"""
        h = np.exp(2*P/1260 * (T_w-T_sat)/2722)
        return h

    def calculate_radial_heat_conduction(self, i, j):
        """计算径向热传导项"""
        # 简化处理,假设均匀网格
        dr = self.channel_diameter / 10  # 径向网格间距
        r = np.linspace(dr/2, self.channel_diameter/2, 10)
        
        # 计算径向温度梯度
        d2T_dr2 = np.zeros_like(r)
        for k in range(1, len(r)-1):
            d2T_dr2[k] = (1/r[k]) * (
                (r[k+1]*self.T_c[i,j] - r[k]*self.T_c[i,j])/dr -
                (r[k]*self.T_c[i,j] - r[k-1]*self.T_c[i,j])/dr
            )/dr
            
        return np.mean(d2T_dr2)  # 返回平均值

    def calculate_fuel_heat_generation(self, i, j):
        """计算燃料热生成"""
        # 简化处理,假设轴向余弦分布
        q_max = 1e6  # 最大线功率密度 (W/m)
        z = i * self.dz
        H = self.channel_length
        q = q_max * np.cos(np.pi * z / H)
        return q / (np.pi * self.channel_diameter)

    def run_transient_simulation(self):
        """运行瞬态计算"""
        print("开始瞬态计算...")
        
        n_time_steps = int(self.total_time / self.time_step)
        
        for time_step in range(n_time_steps):
            # 保存上一时间步的值
            self.save_previous_step_values()
            
            # 在当前时间步内进行迭代求解
            converged, n_iter, final_residual = self.solve_current_time_step()
            
            # 记录残差和迭代次数
            self.history['residuals'].append(final_residual)
            self.history['iterations'].append(n_iter)
            
            # 保存结果
            self.save_timestep_results(time_step)
            
            if not converged:
                print(f"警告：时间步 {time_step + 1} 未达到收敛")
                
            # 更新当前时间
            self.current_time += self.time_step
            
            # 打印计算进度
            if (time_step + 1) % 100 == 0:
                print(f"计算进度: {(time_step + 1)/n_time_steps*100:.1f}% "
                      f"(当前时间: {self.current_time:.3f} s)")
        
        print("瞬态计算完成")
        print(f"结果已保存在 '{self.results_folder}' 文件夹中")
        
    def save_previous_step_values(self):
        """保存上一时间步的场值"""
        self.prev_void_fraction = np.copy(self.void_fraction)
        self.prev_liquid_velocity = np.copy(self.liquid_velocity)
        self.prev_vapor_velocity = np.copy(self.vapor_velocity)
        self.prev_pressure = np.copy(self.pressure)
        self.prev_temperature = np.copy(self.temperature)

    def solve_current_time_step(self):
        """在当前时间步内进行迭代求解"""
        residuals = []
        
        for iteration in range(self.max_iterations):
            # 保存当前迭代步的值用于收敛性检查
            old_void = np.copy(self.void_fraction)
            old_vl = np.copy(self.liquid_velocity)
            old_vv = np.copy(self.vapor_velocity)
            old_p = np.copy(self.pressure)
            old_t = np.copy(self.temperature)
            
            # 求解质量和动量方程
            self.solve_mass_momentum_equations()
            
            # 求解能量方程
            self.solve_energy_equations()
            
            # 检查收敛性
            max_error = self.check_convergence(
                old_void, old_vl, old_vv, old_p, old_t)
            
            residuals.append(max_error)
            
            if max_error < self.convergence_criteria:
                return True, iteration + 1, max_error
                
        return False, self.max_iterations, residuals[-1]
    
    def save_timestep_results(self, time_step):
        """保存当前时间步的结果"""
        try:
            if time_step % self.save_interval == 0:
                # 1. 更新时间相关数据
                self.history['time'].append(self.current_time)
                self.history['time_steps'].append(time_step)
                
                # 2. 更新场变量历史
                self.history['void_fraction_history'].append(np.copy(self.void_fraction))
                self.history['pressure_history'].append(np.copy(self.pressure))
                self.history['temperature_history'].append(np.copy(self.temperature))
                self.history['liquid_velocity_history'].append(np.copy(self.liquid_velocity))
                self.history['vapor_velocity_history'].append(np.copy(self.vapor_velocity))
                self.history['liquid_temp_history'].append(np.copy(self.T_l))
                self.history['vapor_temp_history'].append(np.copy(self.T_v))
                self.history['wall_temp_history'].append(np.copy(self.T_c))
                self.history['heat_flux_liquid_history'].append(np.copy(self.q_wl))
                self.history['heat_flux_vapor_history'].append(np.copy(self.q_wv))
                self.history['phase_change_rate_history'].append(np.copy(self.gamma))
                
                # 3. 更新统计量历史
                self.history['max_void_fraction'].append(np.max(self.void_fraction))
                self.history['avg_pressure'].append(np.mean(self.pressure))
                self.history['max_temperature'].append(np.max(self.temperature))
                self.history['max_liquid_velocity'].append(np.max(self.liquid_velocity))
                self.history['max_vapor_velocity'].append(np.max(self.vapor_velocity))
                
                # 4. 保存CSV文件
                self.save_to_csv(time_step)
                
                # 5. 定期保存完整状态（每100个时间步）
                if time_step % (self.save_interval * 100) == 0:
                    self.save_full_state(time_step)
                    
        except Exception as e:
            print(f"保存时间步 {time_step} 的结果时出错: {str(e)}")
            print(f"错误类型: {type(e).__name__}")
            print(f"错误位置: {traceback.format_exc()}")
            
    def save_to_csv(self, time_step):
        """保存当前时间步的数据到CSV文件"""
        try:
            filename = os.path.join(self.results_folder, f"timestep_{time_step:06d}.csv")
            
            # 使用UTF-8编码保存
            with open(filename, 'w', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)
                
                # 1. 写入时间步信息
                writer.writerow(['时间步信息'])
                writer.writerow(['当前时间 (s)', f"{self.current_time:.6f}"])
                writer.writerow(['时间步编号', time_step])
                writer.writerow([])
                
                # 2. 写入全局参数
                writer.writerow(['全局参数'])
                writer.writerow(['平均压力 (Pa)', f"{np.mean(self.pressure):.2f}"])
                writer.writerow(['最大温度 (K)', f"{np.max(self.temperature):.2f}"])
                writer.writerow(['最大空泡率', f"{np.max(self.void_fraction):.4f}"])
                writer.writerow([])
                
                # 3. 写入详细数据
                headers = [
                    '轴向位置 (m)',
                    '子通道编号',
                    '空泡率',
                    '压力 (Pa)',
                    '液相温度 (K)',
                    '气相温度 (K)',
                    '壁面温度 (K)',
                    '液相速度 (m/s)',
                    '气相速度 (m/s)',
                    '液相壁面热流密度 (W/m²)',
                    '气相壁面热流密度 (W/m²)',
                    '相变率 (kg/m³·s)'
                ]
                writer.writerow(headers)
                
                # 写入每个网格点的数据
                for i in range(self.n_axial):
                    for j in range(self.n_channels):
                        row_data = [
                            f"{i * self.dz:.6f}",
                            j + 1,
                            f"{self.void_fraction[i, j]:.6f}",
                            f"{self.pressure[i, j]:.2f}",
                            f"{self.T_l[i, j]:.2f}",
                            f"{self.T_v[i, j]:.2f}",
                            f"{self.T_c[i, j]:.2f}",
                            f"{self.liquid_velocity[i, j]:.4f}",
                            f"{self.vapor_velocity[i, j]:.4f}",
                            f"{self.q_wl[i, j]:.2f}",
                            f"{self.q_wv[i, j]:.2f}",
                            f"{self.gamma[i, j]:.6f}"
                        ]
                        writer.writerow(row_data)
                        
            print(f"结果已保存到: {filename}")
            
        except Exception as e:
            print(f"保存CSV文件时出错: {str(e)}")
            
    def save_full_state(self, time_step):
        """保存完整的计算状态"""
        try:
            filename = os.path.join(self.results_folder, f"full_state_{time_step:06d}.npz")
            np.savez_compressed(
                filename,
                time=self.current_time,
                time_step=time_step,
                void_fraction=self.void_fraction,
                pressure=self.pressure,
                temperature=self.temperature,
                liquid_velocity=self.liquid_velocity,
                vapor_velocity=self.vapor_velocity,
                T_l=self.T_l,
                T_v=self.T_v,
                T_c=self.T_c,
                q_wl=self.q_wl,
                q_wv=self.q_wv,
                gamma=self.gamma
            )
            print(f"完整状态已保存到: {filename}")
            
        except Exception as e:
            print(f"保存完整状态时出错: {str(e)}")

    def plot_transient_results(self):
        """绘制瞬态计算结果"""
        try:
            # 创建一个包含4个子图的图形
            fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
            
            # 1. 残差历史
            ax1.semilogy(self.history['residuals'])
            ax1.set_title('残差历史')
            ax1.set_xlabel('时间步')
            ax1.set_ylabel('残差')
            ax1.grid(True)
            
            # 2. 最大空泡率随时间的变化
            ax2.plot(self.history['time'], self.history['max_void_fraction'])
            ax2.set_title('最大空泡率随时间的变化')
            ax2.set_xlabel('时间 (s)')
            ax2.set_ylabel('最大空泡率')
            ax2.grid(True)
            
            # 3. 平均压力随时间的变化
            ax3.plot(self.history['time'], self.history['avg_pressure'])
            ax3.set_title('平均压力随时间的变化')
            ax3.set_xlabel('时间 (s)')
            ax3.set_ylabel('压力 (Pa)')
            ax3.grid(True)
            
            # 4. 最大温度随时间的变化
            ax4.plot(self.history['time'], self.history['max_temperature'])
            ax4.set_title('最大温度随时间的变化')
            ax4.set_xlabel('时间 (s)')
            ax4.set_ylabel('温度 (K)')
            ax4.grid(True)
            
            plt.tight_layout()
            plt.show()
            
            # 创建动画展示场分布随时间的变化
            self.animate_field_evolution()
            
        except Exception as e:
            print(f"绘制瞬态结果时出错: {str(e)}")

    def animate_field_evolution(self):
        """创建场分布随时间演化的动画并保存"""
        try:
            # 创建动画图形
            fig = plt.figure(figsize=(15, 5))
            
            # 创建三个子图：空泡率、温度和压力
            ax1 = fig.add_subplot(131, projection='3d')
            ax2 = fig.add_subplot(132, projection='3d')
            ax3 = fig.add_subplot(133, projection='3d')
            
            x = np.arange(self.n_axial)
            y = np.arange(self.n_channels)
            X, Y = np.meshgrid(x, y)
            
            # 确保有足够的历史数据
            n_frames = min(len(self.history['time']),
                          len(self.history['void_fraction_history']),
                          len(self.history['temperature_history']),
                          len(self.history['pressure_history']))
            
            if n_frames == 0:
                print("警告：没有可用的历史数据来创建动画")
                return
                
            def update(frame):
                # 清除当前帧
                ax1.clear()
                ax2.clear()
                ax3.clear()
                
                # 确保frame不会超出数组范围
                frame = min(frame, n_frames - 1)
                
                try:
                    # 绘制空泡率分布
                    surf1 = ax1.plot_surface(X, Y, 
                        self.history['void_fraction_history'][frame].T,
                        cmap=cm.viridis)
                    ax1.set_title(f'空泡率分布\n时间: {self.history["time"][frame]:.3f} s')
                    
                    # 绘制温度分布
                    surf2 = ax2.plot_surface(X, Y,
                        self.history['temperature_history'][frame].T,
                        cmap=cm.viridis)
                    ax2.set_title('温度分布 (K)')
                    
                    # 绘制压力分布
                    surf3 = ax3.plot_surface(X, Y,
                        self.history['pressure_history'][frame].T,
                        cmap=cm.viridis)
                    ax3.set_title('压力分布 (Pa)')
                    
                    # 设置视角
                    for ax in [ax1, ax2, ax3]:
                        ax.view_init(elev=20, azim=45)
                        ax.grid(True)
                        ax.set_xlabel('轴向位置')
                        ax.set_ylabel('子通道编号')
                        
                except IndexError as e:
                    print(f"警告：动画帧 {frame} 超出数据范围")
                    return
            
            # 创建动画
            anim = FuncAnimation(fig, update, 
                                frames=n_frames,
                                interval=200)
            
            # 保存动画
            try:
                anim.save(f'{self.results_folder}/evolution.mp4', 
                         writer='ffmpeg', fps=10)
                print(f"动画已保存到 {self.results_folder}/evolution.mp4")
            except Exception as e:
                print(f"保存动画失败: {str(e)}")
                print("尝试使用其他格式保存...")
                try:
                    anim.save(f'{self.results_folder}/evolution.gif', 
                             writer='pillow', fps=10)
                    print(f"动画已保存为GIF格式")
                except Exception as e:
                    print(f"保存GIF也失败: {str(e)}")
            
            plt.tight_layout()
            plt.show()
            
        except Exception as e:
            print(f"创建动画时出错: {str(e)}")

    def create_results_folder(self):
        """创建结果保存文件夹"""
        if not os.path.exists(self.results_folder):
            os.makedirs(self.results_folder)

    def solve_mass_momentum_equations(self):
        """求解质量和动量方程 - Numba优化版本"""
        # 使用Numba优化的网格点计算
        new_void, new_vl, new_vv = compute_grid_points_numba(
            self.n_axial, self.n_channels,
            self.void_fraction, self.prev_void_fraction,
            self.vapor_velocity, self.liquid_velocity,
            self.prev_liquid_velocity, self.prev_vapor_velocity,
            self.pressure, self.temperature,
            self.time_step, self.dz
        )
        
        # 更新结果
        self.void_fraction = new_void
        self.liquid_velocity = new_vl
        self.vapor_velocity = new_vv
        
    def solve_energy_equations(self):
        """求解能量方程 - Numba优化��本"""
        # 使用Numba优化的能量点计算
        new_temp = compute_energy_points_numba(
            self.n_axial, self.n_channels,
            self.temperature, self.prev_temperature,
            self.liquid_velocity, self.pressure,
            self.time_step, self.dz
        )
        
        # 更新温度
        self.temperature = new_temp
        
    def calculate_wall_friction(self, i, j, phase, direction='axial'):
        """计算壁面摩擦力 - Numba优化版本"""
        if phase == 'liquid':
            props = self.get_fluid_properties(self.pressure[i,j], self.T_l[i,j], 'liquid')
            if direction == 'axial':
                velocity = self.liquid_velocity[i,j]
            else:
                velocity = self.v_l[i,j]
        else:
            props = self.get_fluid_properties(self.pressure[i,j], self.T_v[i,j], 'vapor')
            if direction == 'axial':
                velocity = self.vapor_velocity[i,j]
            else:
                velocity = self.v_v[i,j]
                
        if props is None:
            return 0.0
            
        return calculate_wall_friction_numba(
            velocity, self.channel_diameter,
            props['density'], props['viscosity']
        )
        
    def calculate_residual(self):
        """计算残差 - Numba优化版本"""
        return calculate_residual_numba(
            self.n_axial, self.n_channels,
            self.liquid_velocity, self.liquid_enthalpy,
            self.dz
        )

    def get_latent_heat(self, pressure):
        """获取给定压力下的汽化潜热"""
        try:
            h_f = PropsSI('H', 'P', pressure, 'Q', 0, 'Water')  # 饱和液焓
            h_g = PropsSI('H', 'P', pressure, 'Q', 1, 'Water')  # 饱和气焓
            return h_g - h_f
        except:
            return 2.257e6  # 返回1个大气压下的近似值

    def get_saturation_temperature(self, pressure):
        """
        计算给定压力下的饱和温度
        
        参数：
        pressure: float, 压力 (Pa)
        
        返回：
        float: 饱和温度 (K)
        """
        try:
            return PropsSI('T', 'P', pressure, 'Q', 0, 'Water')
        except Exception as e:
            print(f"警告：饱和温度计算失败 (P={pressure})")
            print(f"错误信息: {str(e)}")
            return 373.15  # 返回1个大气压下的近似值

    def calculate_single_phase_htc(self, i, j, phase):
        """
        计算单相对流换热系数
        
        参数：
        i, j: int, 网格位置
        phase: str, 'liquid' 或 'vapor'
        
        返回：
        float: 换热系数 (W/m²·K)
        """
        if phase == 'liquid':
            props = self.get_fluid_properties(self.pressure[i,j], self.T_l[i,j], 'liquid')
            velocity = self.liquid_velocity[i,j]
        else:
            props = self.get_fluid_properties(self.pressure[i,j], self.T_v[i,j], 'vapor')
            velocity = self.vapor_velocity[i,j]
        
        if props is None:
            return 0.0
        
        Re = props['density'] * velocity * self.channel_diameter / props['viscosity']
        Pr = props['viscosity'] * props.get('specific_heat', 4200) / props.get('thermal_conductivity', 0.6)
        
        # Dittus-Boelter correlation
        if Re > 10000 and 0.7 <= Pr <= 160:
            Nu = 0.023 * Re**0.8 * Pr**0.4
        else:
            Nu = 4.36  # 层流或过渡区域的简化处理
        
        h = Nu * props.get('thermal_conductivity', 0.6) / self.channel_diameter
        return h

    def check_convergence(self, old_void, old_vl, old_vv, old_p, old_t):
        """检查收敛性"""
        max_error = max(
            np.max(np.abs(self.void_fraction - old_void)),
            np.max(np.abs(self.liquid_velocity - old_vl)),
            np.max(np.abs(self.vapor_velocity - old_vv)),
            np.max(np.abs(self.pressure - old_p)) / self.inlet_pressure,
            np.max(np.abs(self.temperature - old_t)) / self.inlet_temp
        )
        return max_error

    def setup_staggered_grid(self):
        """设置交错网格"""
        # 标量网格（压力、温度、空泡率等）
        self.scalar_grid = np.zeros((self.n_axial, self.n_channels))
        
        # 轴向速度网格（错开半个网格）
        self.axial_velocity_grid = np.zeros((self.n_axial + 1, self.n_channels))
        
        # 横向速度网格（错开半个网格）
        self.lateral_velocity_grid = np.zeros((self.n_axial, self.n_channels + 1))
        
        # 网格中心点坐标
        self.z_centers = np.linspace(0, self.channel_length, self.n_axial)
        self.z_faces = np.linspace(-self.dz/2, self.channel_length + self.dz/2, self.n_axial + 1)
        
    def build_preconditioner(self):
        """构建基于物理过程的预处理器"""
        n = self.n_axial * self.n_channels
        M = lil_matrix((n, n))
        
        # 构建预处理器矩阵（仅考虑轴向相邻单元）
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                idx = i * self.n_channels + j
                
                # 中心单元
                M[idx, idx] = 1.0
                
                # 轴向相邻单元
                if i > 0:
                    M[idx, (i-1) * self.n_channels + j] = -0.5
                if i < self.n_axial - 1:
                    M[idx, (i+1) * self.n_channels + j] = -0.5
                    
        return M.tocsr()
        
    def solve_pkm(self):
        """使用Picard Krylov方法求解方程组"""
        converged = False
        picard_iter = 0
        
        while not converged and picard_iter < self.max_picard_iterations:
            # 1. 计算热通量
            self.calculate_heat_flux()
            
            # 2. 求解动量方程
            self.solve_momentum_equations()
            
            # 3. 构建质量-能量方程组
            A, b = self.build_mass_energy_system()
            
            # 4. 构建预处理器
            M = self.build_preconditioner()
            
            # 5. 使用PGMRES求解线性方程组
            try:
                x, info = gmres(A, b, M=M, tol=self.krylov_tolerance, 
                              maxiter=self.max_krylov_iterations)
                
                if info > 0:
                    print(f"警告：GMRES在{info}次迭代后未收敛")
                    
                # 更新解
                self.update_solution(x)
                
                # 计算收敛性
                residual = self.calculate_residual()
                converged = residual < self.picard_tolerance
                
            except Exception as e:
                print(f"错误：GMRES求解失败 - {str(e)}")
                # 如果求解失败，减小时间步长重试
                self.time_step *= 0.5
                if self.time_step < 1e-6:
                    raise Exception("时间步长过小，计算终止")
                continue
                
            picard_iter += 1
            
        return converged, picard_iter
        
    def build_mass_energy_system(self):
        """构建质量-能量方程组的系数矩阵和右端向量"""
        n = self.n_axial * self.n_channels
        A = lil_matrix((n, n))
        b = np.zeros(n)
        
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                idx = i * self.n_channels + j
                
                # 1. 质量守恒方程
                # 轴向质量流量
                if i > 0:
                    A[idx, idx-self.n_channels] = -self.liquid_velocity[i-1,j]
                A[idx, idx] = self.liquid_velocity[i,j]
                if i < self.n_axial - 1:
                    A[idx, idx+self.n_channels] = -self.liquid_velocity[i+1,j]
                
                # 横向质量流量
                for k in self.channel_connections[j+1]:
                    k_idx = i * self.n_channels + (k-1)
                    w_jk = self.calculate_cross_flow(i, j+1, k)
                    A[idx, k_idx] = w_jk
                    
                # 2. 能量守恒方程
                # 轴向能量输运
                if i > 0:
                    A[idx, idx-self.n_channels] += -self.liquid_velocity[i-1,j] * self.liquid_enthalpy[i-1,j]
                A[idx, idx] += self.liquid_velocity[i,j] * self.liquid_enthalpy[i,j]
                if i < self.n_axial - 1:
                    A[idx, idx+self.n_channels] += -self.liquid_velocity[i+1,j] * self.liquid_enthalpy[i+1,j]
                
                # 加热项
                q_wall = self.calculate_wall_heat_flux(i, j, 'liquid')
                b[idx] = q_wall
                
        return A.tocsr(), b
        
    def solve_momentum_equations(self):
        """求解动量方程"""
        # 轴向动量方程
        for i in range(1, self.n_axial):
            for j in range(self.n_channels):
                # 获取交错网格上的速度
                v_up = self.axial_velocity_grid[i-1, j]
                v_center = self.axial_velocity_grid[i, j]
                v_down = self.axial_velocity_grid[i+1, j] if i < self.n_axial-1 else v_center
                
                # 计算压力梯度
                dp_dz = (self.pressure[i,j] - self.pressure[i-1,j]) / self.dz
                
                # 计算摩擦力
                f_wall = self.calculate_wall_friction(i, j, 'liquid')
                
                # 更新轴向速度
                self.axial_velocity_grid[i,j] = v_center - self.time_step * (
                    v_center * (v_down - v_up)/(2*self.dz) +  # 对流项
                    dp_dz/self.get_fluid_properties(self.pressure[i,j], 
                        self.temperature[i,j], 'liquid')['density'] +  # 压力项
                    f_wall +  # 摩擦力
                    9.81  # 重力项
                )
        
        # 横向动量方程
        for i in range(self.n_axial):
            for j in range(self.n_channels-1):
                # 计算横向压力梯度
                dp_dy = self.calculate_lateral_pressure_gradient(i, j)
                
                # 计算横向摩擦力
                f_lat = self.calculate_wall_friction(i, j, 'liquid', 'lateral')
                
                # 更新横向速度
                self.lateral_velocity_grid[i,j] = self.lateral_velocity_grid[i,j] - self.time_step * (
                    dp_dy/self.get_fluid_properties(self.pressure[i,j], 
                        self.temperature[i,j], 'liquid')['density'] +
                    f_lat
                )
                
    def update_solution(self, x):
        """更新解向量"""
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                idx = i * self.n_channels + j
                
                # 更新压力
                self.pressure[i,j] = x[idx]
                
                # 更新温度（通过压力和焓计算）
                h = x[idx + self.n_axial * self.n_channels]
                self.temperature[i,j] = self.calculate_temperature_from_enthalpy(
                    h, self.pressure[i,j], 'liquid')
                
    def calculate_heat_flux(self):
        """计算热通量"""
        # 计算热通量
        # ...

@njit(parallel=True)
def compute_grid_points_numba(n_axial, n_channels, void_fraction, prev_void_fraction, 
                            vapor_velocity, liquid_velocity, prev_liquid_velocity,
                            prev_vapor_velocity, pressure, temperature, time_step, dz):
    """使Numba优化的网格点计算"""
    new_void = np.zeros_like(void_fraction)
    new_vl = np.zeros_like(liquid_velocity)
    new_vv = np.zeros_like(vapor_velocity)
    
    for i in prange(1, n_axial):
        for j in prange(n_channels):
            # 计算空间导数项
            if i > 0:
                d_alpha_dx = (void_fraction[i,j] - void_fraction[i-1,j]) / dz
                d_vl_dx = (liquid_velocity[i,j] - liquid_velocity[i-1,j]) / dz
                d_vv_dx = (vapor_velocity[i,j] - vapor_velocity[i-1,j]) / dz
                dp_dx = (pressure[i,j] - pressure[i-1,j]) / dz
            else:
                d_alpha_dx = 0
                d_vl_dx = 0
                d_vv_dx = 0
                dp_dx = 0
            
            # 更新值
            new_void[i,j] = prev_void_fraction[i,j] - time_step * (
                vapor_velocity[i,j] * d_alpha_dx
            )
            
            new_vl[i,j] = prev_liquid_velocity[i,j] - time_step * (
                liquid_velocity[i,j] * d_vl_dx +
                dp_dx +
                9.81
            )
            
            new_vv[i,j] = prev_vapor_velocity[i,j] - time_step * (
                vapor_velocity[i,j] * d_vv_dx +
                dp_dx +
                9.81
            )
            
            # 限制物理范围
            new_void[i,j] = min(max(new_void[i,j], 0.0), 0.9999)
            new_vl[i,j] = max(new_vl[i,j], 0.0)
            new_vv[i,j] = max(new_vv[i,j], 0.0)
            
    return new_void, new_vl, new_vv

@njit(parallel=True)
def compute_energy_points_numba(n_axial, n_channels, temperature, prev_temperature,
                              liquid_velocity, pressure, time_step, dz):
    """使用Numba优化的能量点计算"""
    new_temp = np.zeros_like(temperature)
    
    for i in prange(1, n_axial):
        for j in prange(n_channels):
            # 计算空间导数项
            if i > 0:
                d_T_dx = (temperature[i,j] - temperature[i-1,j]) / dz
            else:
                d_T_dx = 0
            
            # 更新温度
            new_temp[i,j] = prev_temperature[i,j] - time_step * (
                liquid_velocity[i,j] * d_T_dx
            )
            
    return new_temp

@njit
def calculate_wall_friction_numba(velocity, diameter, density, viscosity):
    """使用Numba优化的壁面摩擦力计算"""
    # 计算Reynolds数
    Re = density * abs(velocity) * diameter / viscosity
    
    # 计算摩擦系数
    if Re < 2300:  # 层流
        f = 64 / Re
    else:  # 湍流
        f = 0.316 * Re**(-0.25)
        
    # 计算壁面摩擦力
    tau_w = 0.5 * f * density * velocity * abs(velocity)
    
    return tau_w

@njit(parallel=True)
def calculate_residual_numba(n_axial, n_channels, liquid_velocity, liquid_enthalpy, dz):
    """使用Numba优化的残差计算"""
    mass_residual = np.zeros((n_axial, n_channels))
    energy_residual = np.zeros((n_axial, n_channels))
    
    for i in prange(1, n_axial):
        for j in prange(n_channels):
            # 轴向质量流量差
            dm_dz = (liquid_velocity[i,j] - liquid_velocity[i-1,j]) / dz
            mass_residual[i,j] = dm_dz
            
            # 轴向能量输运
            dh_dz = (liquid_enthalpy[i,j] - liquid_enthalpy[i-1,j]) / dz
            energy_residual[i,j] = dh_dz
            
    return max(np.max(np.abs(mass_residual)), np.max(np.abs(energy_residual)))

@njit(parallel=True)
def solve_two_fluid_equations_numba(n_axial, n_channels, 
                                  alpha_l, alpha_v, rho_l, rho_v,
                                  u_l, u_v, v_l, v_v, w_l, w_v,
                                  P, T_l, T_v, T_c, h_l, h_v,
                                  dx, dy, dz, dt):
    """使用Numba优化的双流体模型求解器"""
    # 计算CFL时间步长
    dt_cfl = calculate_cfl_timestep(u_l, u_v, v_l, v_v, dx, dy, dz)
    
    # 使用较小的时间步长
    dt = min(dt, dt_cfl)
    
    # 初始化新的场值
    new_alpha_l = np.zeros_like(alpha_l)
    new_alpha_v = np.zeros_like(alpha_v)
    new_u_l = np.zeros_like(u_l)
    new_u_v = np.zeros_like(u_v)
    new_v_l = np.zeros_like(v_l)
    new_v_v = np.zeros_like(v_v)
    new_h_l = np.zeros_like(h_l)
    new_h_v = np.zeros_like(h_v)
    
    for i in prange(1, n_axial-1):
        for j in prange(1, n_channels-1):
            # ====== 1. 计算相变质量流量 ======
            T_sat = calculate_saturation_temperature_numba(P[i,j])
            
            # 过冷沸腾
            if T_l[i,j] < T_sat and T_c[i,j] > T_sat:
                h_sc = calculate_subcooled_boiling_htc_numba(T_c[i,j], T_sat, T_l[i,j], P[i,j])
                Q_sc = h_sc * (T_c[i,j] - T_l[i,j])
                gamma_sc = Q_sc / calculate_latent_heat_numba(P[i,j])
            else:
                gamma_sc = 0.0
                
            # 饱和核态沸腾
            if T_l[i,j] >= T_sat and T_c[i,j] > T_sat:
                h_sh = calculate_nucleate_boiling_htc_numba(T_c[i,j], T_sat, P[i,j])
                Q_sh = h_sh * (T_c[i,j] - T_sat)
                gamma_sh = Q_sh / calculate_latent_heat_numba(P[i,j])
            else:
                gamma_sh = 0.0
                
            # 净相变质量流量
            gamma = gamma_sh - gamma_sc
            
            # ====== 2. 计算湍流混合项 ======
            # 液相湍流混合
            T_ml = calculate_turbulent_mixing_numba(alpha_l[i,j], alpha_l[i,j+1], 
                                                  rho_l[i,j], dy)
            # 气相湍流混合
            T_mv = calculate_turbulent_mixing_numba(alpha_v[i,j], alpha_v[i,j+1], 
                                                  rho_v[i,j], dy)
            
            # ====== 3. 计算界面力 ======
            F_lv = calculate_interfacial_force_numba(alpha_l[i,j], rho_l[i,j],
                                                   u_l[i,j], u_v[i,j],
                                                   v_l[i,j], v_v[i,j])
            
            # ====== 4. 液相质量守恒方程 ======
            d_alpha_l_dt = -(1/rho_l[i,j]) * (
                # 轴向质量通量导数
                (alpha_l[i+1,j]*rho_l[i+1,j]*u_l[i+1,j] - 
                 alpha_l[i-1,j]*rho_l[i-1,j]*u_l[i-1,j])/(2*dz) +
                # 横向质量通量导数
                (alpha_l[i,j+1]*rho_l[i,j+1]*v_l[i,j+1] - 
                 alpha_l[i,j-1]*rho_l[i,j-1]*v_l[i,j-1])/(2*dy)
            ) - gamma/rho_l[i,j] + T_ml
            
            # ====== 5. 气相质量守恒方程 ======
            d_alpha_v_dt = -(1/rho_v[i,j]) * (
                # 轴向质量通量导数
                (alpha_v[i+1,j]*rho_v[i+1,j]*u_v[i+1,j] - 
                 alpha_v[i-1,j]*rho_v[i-1,j]*u_v[i-1,j])/(2*dz) +
                # 横向质量通量导数
                (alpha_v[i,j+1]*rho_v[i,j+1]*v_v[i,j+1] - 
                 alpha_v[i,j-1]*rho_v[i,j-1]*v_v[i,j-1])/(2*dy)
            ) + gamma/rho_v[i,j] + T_mv
            
            # ====== 6. 液相轴向动量方程 ======
            # 壁面摩擦力
            F_wl = calculate_wall_friction_numba(u_l[i,j], 0.01, rho_l[i,j], 1e-3)
            
            d_ul_dt = -(
                # 对流项
                u_l[i,j] * (u_l[i+1,j] - u_l[i-1,j])/(2*dz) +
                v_l[i,j] * (u_l[i,j+1] - u_l[i,j-1])/(2*dy) +
                # 压力梯度
                (1/rho_l[i,j]) * (P[i+1,j] - P[i-1,j])/(2*dz) +
                # 重力项
                9.81 +
                # 壁面摩擦力
                F_wl/rho_l[i,j] -
                # 界面力
                F_lv/rho_l[i,j] -
                # 相变动量
                gamma * u_l[i,j]/rho_l[i,j]
            )
            
            # ====== 7. 气相轴向动量方程 ======
            # 壁面摩擦力
            F_wv = calculate_wall_friction_numba(u_v[i,j], 0.01, rho_v[i,j], 1e-5)
            
            d_uv_dt = -(
                # 对流项
                u_v[i,j] * (u_v[i+1,j] - u_v[i-1,j])/(2*dz) +
                v_v[i,j] * (u_v[i,j+1] - u_v[i,j-1])/(2*dy) +
                # 压力梯度
                (1/rho_v[i,j]) * (P[i+1,j] - P[i-1,j])/(2*dz) +
                # 重力项
                9.81 +
                # 壁面摩擦力
                F_wv/rho_v[i,j] +
                # 界面力
                F_lv/rho_v[i,j] +
                # 相变动量
                gamma * u_v[i,j]/rho_v[i,j]
            )
            
            # ====== 8. 液相横向动量方程 ======
            d_vl_dt = -(
                # 对流项
                u_l[i,j] * (v_l[i+1,j] - v_l[i-1,j])/(2*dz) +
                v_l[i,j] * (v_l[i,j+1] - v_l[i,j-1])/(2*dy) -
                # 界面力
                F_lv/rho_l[i,j] -
                # 相变动量
                gamma * v_l[i,j]/rho_l[i,j]
            )
            
            # ====== 9. 气相横向动量方程 ======
            d_vv_dt = -(
                # 对流项
                u_v[i,j] * (v_v[i+1,j] - v_v[i-1,j])/(2*dz) +
                v_v[i,j] * (v_v[i,j+1] - v_v[i,j-1])/(2*dy) +
                # 压力梯度
                (1/rho_v[i,j]) * (P[i,j+1] - P[i,j-1])/(2*dy) +
                # 界面力
                F_lv/rho_v[i,j] +
                # 相变动量
                gamma * v_v[i,j]/rho_v[i,j]
            )
            
            # ====== 10. 液相能量方程 - 添加燃料热生成和径向传导 ======
            # 壁面换热
            q_wl = calculate_wall_heat_flux_numba(T_c[i,j], T_l[i,j], h_l[i,j])
            
            # 计算燃料热生成
            q_fuel = calculate_fuel_heat_generation_numba(i, n_axial, dz)
            
            # 计算径向热传导
            q_rad = calculate_radial_heat_conduction_numba(T_c[i,j], T_l[i,j], dx)
            
            d_hl_dt = -(
                # 对流项
                u_l[i,j] * (h_l[i+1,j] - h_l[i-1,j])/(2*dz) +
                v_l[i,j] * (h_l[i,j+1] - h_l[i,j-1])/(2*dy) +
                # 压力功
                alpha_l[i,j] * (P[i,j] - P[i-1,j])/dt +
                # 壁面换热
                q_wl/(rho_l[i,j] * alpha_l[i,j]) +
                # 相变热量
                gamma * h_l[i,j]/(rho_l[i,j] * alpha_l[i,j]) +
                # 添加燃料热生成和径向传导
                (q_fuel + q_rad)/(rho_l[i,j] * alpha_l[i,j])
            )
            
            # ====== 11. 气相能量方程 - 添加燃料热生成和径向传导 ======
            # 壁面换热
            q_wv = calculate_wall_heat_flux_numba(T_c[i,j], T_v[i,j], h_v[i,j])
            
            d_hv_dt = -(
                # 对流项
                u_v[i,j] * (h_v[i+1,j] - h_v[i-1,j])/(2*dz) +
                v_v[i,j] * (h_v[i,j+1] - h_v[i,j-1])/(2*dy) +
                # 压力功
                alpha_v[i,j] * (P[i,j] - P[i-1,j])/dt +
                # 壁面换热
                q_wv/(rho_v[i,j] * alpha_v[i,j]) +
                # 相变热量
                gamma * h_v[i,j]/(rho_v[i,j] * alpha_v[i,j]) +
                # 添加燃料热生成和径向传导（气相部分）
                (q_fuel * alpha_v[i,j] + q_rad)/(rho_v[i,j] * alpha_v[i,j])
            )
            
            # ====== 12. 时间推进 ======
            new_alpha_l[i,j] = alpha_l[i,j] + dt * d_alpha_l_dt
            new_alpha_v[i,j] = alpha_v[i,j] + dt * d_alpha_v_dt
            new_u_l[i,j] = u_l[i,j] + dt * d_ul_dt
            new_u_v[i,j] = u_v[i,j] + dt * d_uv_dt
            new_v_l[i,j] = v_l[i,j] + dt * d_vl_dt
            new_v_v[i,j] = v_v[i,j] + dt * d_vv_dt
            new_h_l[i,j] = h_l[i,j] + dt * d_hl_dt
            new_h_v[i,j] = h_v[i,j] + dt * d_hv_dt
            
            # 限制物理范围
            new_alpha_l[i,j] = min(max(new_alpha_l[i,j], 0.0), 1.0)
            new_alpha_v[i,j] = 1.0 - new_alpha_l[i,j]
            
    # 在返回结果之前检查数值稳定性
    dt_new, needs_restart = check_numerical_stability(
        new_alpha_l, new_alpha_v, new_u_l, new_u_v, new_v_l, new_v_v,
        P, T_l, T_v, rho_l, rho_v, h_l, h_v, dt
    )
    
    if needs_restart:
        # 如果需要重新计算，返回特殊标记
        return (None, None, None, None, None, None, None, None, dt_new, True)
    
    return (new_alpha_l, new_alpha_v, new_u_l, new_u_v, 
            new_v_l, new_v_v, new_h_l, new_h_v, dt_new, False)

@njit
def calculate_turbulent_mixing_numba(alpha_i, alpha_j, rho, dy):
    """计算湍流混合项"""
    C_e = 0.01  # 涡流扩散率
    L_tur = dy  # 湍流混合长度
    alpha_eq_i = 0.5  # 平衡空泡分数
    
    return C_e * dy * L_tur * rho * (alpha_j - alpha_i - (alpha_j - alpha_eq_i))

@njit
def calculate_interfacial_force_numba(alpha_l, rho_l, u_l, u_v, v_l, v_v):
    """计算界面力"""
    C_i = 0.44  # 界面阻力系数
    d_b = 0.001  # 气泡直径
    
    # 计算相对速度
    u_r = u_v - u_l
    v_r = v_v - v_l
    V_r = np.sqrt(u_r**2 + v_r**2)
    
    # 计算界面面积密度
    a_i = 6 * (1-alpha_l) / d_b
    
    return 0.5 * C_i * rho_l * V_r * V_r * a_i

@njit
def calculate_wall_heat_flux_numba(T_w, T_f, h):
    """计算壁面热流密度"""
    return h * (T_w - T_f)

@njit
def calculate_saturation_temperature_numba(P):
    """计算饱和温度"""
    return 373.15 + 0.0001 * (P - 1.013e5)  # 简化的饱和温度计算

@njit
def calculate_latent_heat_numba(P):
    """计算汽化潜热"""
    return 2.257e6  # 常压下水的汽化潜热

@njit
def calculate_cfl_timestep(u_l, u_v, v_l, v_v, dx, dy, dz):
    """计算基于CFL条件的时��步长"""
    # ��算最大速度
    max_velocity = max(
        np.max(np.abs(u_l)),
        np.max(np.abs(u_v)),
        np.max(np.abs(v_l)),
        np.max(np.abs(v_v))
    )
    
    # CFL条件，取最小网格间距
    min_dx = min(dx, dy, dz)
    
    # CFL数设为0.5以保证稳定性
    cfl = 0.5
    
    # 计算时间步长
    dt = cfl * min_dx / (max_velocity + 1e-10)
    
    return dt

@njit
def check_numerical_stability(alpha_l, alpha_v, u_l, u_v, v_l, v_v, P, T_l, T_v,
                            rho_l, rho_v, h_l, h_v, dt_old):
    """检查数值稳定性并调整时间步长"""
    # 1. 检查物理量的变化率
    max_alpha_change = max(
        np.max(np.abs(np.diff(alpha_l))),
        np.max(np.abs(np.diff(alpha_v)))
    )
    
    max_velocity_change = max(
        np.max(np.abs(np.diff(u_l))),
        np.max(np.abs(np.diff(u_v))),
        np.max(np.abs(np.diff(v_l))),
        np.max(np.abs(np.diff(v_v)))
    )
    
    max_pressure_change = np.max(np.abs(np.diff(P))) / np.mean(P)
    max_temp_change = max(
        np.max(np.abs(np.diff(T_l))) / np.mean(T_l),
        np.max(np.abs(np.diff(T_v))) / np.mean(T_v)
    )
    
    # 2. 设置阈值
    alpha_threshold = 0.1
    velocity_threshold = 0.2
    pressure_threshold = 0.1
    temp_threshold = 0.1
    
    # 3. 根据变化率调整时间步长
    dt_factor = 1.0
    
    if max_alpha_change > alpha_threshold:
        dt_factor = min(dt_factor, 0.5)
    if max_velocity_change > velocity_threshold:
        dt_factor = min(dt_factor, 0.5)
    if max_pressure_change > pressure_threshold:
        dt_factor = min(dt_factor, 0.5)
    if max_temp_change > temp_threshold:
        dt_factor = min(dt_factor, 0.5)
        
    # 如果所有变化率都很小，可以适当增加时间步长
    if (max_alpha_change < alpha_threshold/2 and
        max_velocity_change < velocity_threshold/2 and
        max_pressure_change < pressure_threshold/2 and
        max_temp_change < temp_threshold/2):
        dt_factor = min(1.2, 1.0/dt_factor)
        
    # 4. 检查物理约束
    if (np.any(alpha_l < 0) or np.any(alpha_l > 1) or
        np.any(P < 0) or
        np.any(T_l < 273) or np.any(T_v < 273) or
        np.any(rho_l < 0) or np.any(rho_v < 0)):
        dt_factor = 0.1  # 严重违反物理约束，大幅减小时间步长
        
    # 5. 限制时间步长的变化范围
    dt_new = dt_old * dt_factor
    dt_new = max(1e-6, min(dt_new, 1e-2))  # 限制在合理范围内
    
    return dt_new, dt_factor < 1.0

@njit
def calculate_fuel_heat_generation_numba(i, n_axial, dz):
    """计算燃料热生成 - Numba优化版本"""
    q_max = 1e6  # 最大线功率密度 (W/m)
    z = i * dz
    H = n_axial * dz
    # 使用余弦分布模拟轴向功率分布
    q = q_max * np.cos(np.pi * z / H)
    return q

@njit
def calculate_radial_heat_conduction_numba(T_c, T_f, dr):
    """计算径向热传导 - Numba优化版���"""
    k = 30.0  # 假设的导热系数 (W/m·K)
    # 简化的径向热传导计算
    q_rad = k * (T_c - T_f) / dr
    return q_rad

if __name__ == "__main__":
    # 创建分析实例
    analysis = SubchannelAnalysis(n_axial=1500)
    
    # 运行瞬态模拟
    analysis.run_transient_simulation()
    
    # 显示瞬态结果
    analysis.plot_transient_results()
    
    # 显示最终状态的详细分布
    analysis.plot_results()